package trie

// TrieNode 前缀树
type TrieNode struct {
	//End 是否有单词以该节点结尾（标记是否有这个单词）
	End bool
	//PassCnt 统计有多少个单词经过该节点
	PassCnt int
	//Children 指向子节点
	Children [26]*TrieNode
}

type TrieTree struct {
	Root *TrieNode
}

func NewTrieTree() *TrieTree {
	return &TrieTree{
		Root: &TrieNode{},
	}
}

func (t *TrieTree) Insert(word string) {
	if t.Search(word) {
		return
	}

	move := t.Root
	for _, v := range word {
		idx := v - 'a'
		if move.Children[idx] == nil {
			move.Children[idx] = &TrieNode{
				End:      false,
				PassCnt:  0,
				Children: [26]*TrieNode{},
			}
		}
		move = move.Children[idx]
		move.PassCnt++
	}
	move.End = true
}

func (t *TrieTree) Search(word string) bool {
	node := t.searchNode(word)
	return node != nil && node.End
}

func (t *TrieTree) Delete(word string) bool {
	if !t.Search(word) {
		return false
	}
	move := t.Root
	for _, v := range word {
		idx := v - 'a'
		move.Children[idx].PassCnt--
		if move.Children[idx].PassCnt == 0 {
			move.Children[idx] = nil
			return true
		}
		move = move.Children[idx]
	}

	move.End = false
	return true
}

func (t *TrieTree) PassCnt(word string) int {
	node := t.searchNode(word)
	if node == nil {
		return 0
	}
	return node.PassCnt
}

func (t *TrieTree) searchNode(word string) *TrieNode {
	// 获得root节点
	// 遍历传入word的每个字母
	// 从root节点开始获取
	move := t.Root
	for _, v := range word {
		idx := v - 'a'
		if move.Children[idx] == nil {
			return nil
		}
		move = move.Children[idx]
	}

	return move
}
